#include <stdio.h>
#include <stdint.h>
#include <fstream>
#include <vector>
#include <queue>

void dumpbin(void *buf, size_t sz, const std::string &path) {
  std::ofstream ofs;
  ofs.open(path, std::ios::out | std::ios::trunc | std::ios::binary);
  ofs.write((const char *)buf, sz);
  // printf("# dump %s \n", path.c_str());
}

template<typename T>
void memset(void *buf, const T raw_val, int sz) {
  for(int id = 0; id < sz; ++id){
    T val = raw_val;
    if(id%2 == 0){
      val += id;
    }else{
      val -= id;
    }
    ((T*)buf)[id] = val;
  }
}

void topk_host(
  int m, int k, float *A, float *B) {

  std::priority_queue<float, std::vector<float>, std::greater<float>> q_;

  for(int i = 0; i < m; ++i){
    if(q_.size() < k){
      q_.push(A[i]);
    }else{
      if(A[i] > q_.top()){
        q_.push(A[i]);
        q_.pop();
      }
    }
  }
  for(int i = 0; i < k; ++i){
    B[k-i-1] = q_.top();
    q_.pop();
  }
  // for(int i = 0; i < k; ++i){
  //   printf("%0.2f, ", B[i]);
  // }
}

struct Case{
  int m;
  int k;
  float A_val;
  std::string A_path;
  std::string B_path;
  float latency;
};

#include <chrono>

static std::uint64_t now() {
  return std::chrono::duration_cast<std::chrono::milliseconds>
    (std::chrono::high_resolution_clock::now().time_since_epoch()).count();
}

int main(int argc, char *argv[]) {
  std::vector<Case> cases;
  for(int i = 41; i <= 80; i+=1){
    cases.push_back({1100, i, 0.f, std::string("data/c0_A.bin.")+std::to_string(i), std::string("data/c0_B.bin.")+std::to_string(i)});
    cases.push_back({5028, i, 0.f, std::string("data/c11_A.bin.")+std::to_string(i), std::string("data/c11_B.bin.")+std::to_string(i)});
    cases.push_back({10028, i, 0.f, std::string("data/c12_A.bin.")+std::to_string(i), std::string("data/c12_B.bin.")+std::to_string(i)});
    cases.push_back({50028, i, 0.f, std::string("data/c9_A.bin.")+std::to_string(i), std::string("data/c9_B.bin.")+std::to_string(i)});
    cases.push_back({100028, i, 0.f, std::string("data/c4_A.bin.")+std::to_string(i), std::string("data/c4_B.bin.")+std::to_string(i)});
    cases.push_back({530167, i, 0.f, std::string("data/c6_A.bin.")+std::to_string(i), std::string("data/c6_B.bin.")+std::to_string(i)});
    cases.push_back({1130167, i, 0.f, std::string("data/c8_A.bin.")+std::to_string(i), std::string("data/c8_B.bin.")+std::to_string(i)});
    cases.push_back({5000167, i, 0.f, std::string("data/c14_A.bin.")+std::to_string(i), std::string("data/c14_B.bin.")+std::to_string(i)});
    cases.push_back({10000167, i, 0.f, std::string("data/c15_A.bin.")+std::to_string(i), std::string("data/c15_B.bin.")+std::to_string(i)});
  }

  for(int id = 0; id < cases.size(); ++id){
    auto &c = cases[id];
    const int m = c.m, k = c.k;
    std::string A_path = c.A_path, B_path = c.B_path;
    float A_val = c.A_val;

    float *A = new float[m];
    float *B = new float[k];
    memset(A, A_val, m);
    memset(B, 0, k);
    
    int loop = 1;
    int start = now();
    topk_host(m, k, A, B);
    int end = now();
    c.latency =float(end-start)/float(loop);

    dumpbin(A, m*sizeof(float), A_path);
    dumpbin(B, k*sizeof(float), B_path);
    delete [] A;
    delete [] B;
    printf("{\"m\": %d, \"k\": %d, \"A_path\": \"%s\", \"B_path\": \"%s\", \"latency\": %0.2f}\n", 
      m, k, A_path.c_str(), B_path.c_str(), c.latency);
  }
  return 0;
}